{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 基于AReaL强化学习训练长程搜索智能体（Search Agent）\n",
    "\n",
    "本教程利用AReaL中的基本组件快速搭建一个强化学习流程，用来训练一个可以进行长程搜索的智能体。\n",
    "\n",
    "该教程包括以下步骤：\n",
    "1. 实验准备（包括从yaml加载实验配置，配置环境变量，启动SGLang服务器，启动本地RAG服务器，加载训练数据集）\n",
    "2. 定义简单的工作流，多次调用搜索工具；\n",
    "3. 每次生成**多条**轨迹（i.e., GRPO)；\n",
    "4. 测试工作流；\n",
    "5. 将工作流接入端到端GRPO强化学习训练；"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## 实验准备\n",
    "### 加载实验配置\n",
    "\n",
    "通过`load_expr_config`加载预定义的asearcher基于local RAG训练的yaml实验配置模板。\n",
    "\n",
    "实验配置模板内配置了优化器、模板、学习率等参数，可以直接使用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import asdict, dataclass, field\n",
    "\n",
    "from areal.api.cli_args import GRPOConfig, load_expr_config\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class AgentRLConfig(GRPOConfig):\n",
    "    max_turns: int = field(\n",
    "        default=128, metadata={\"help\": \"maximum number of turns per trajectory\"}\n",
    "    )\n",
    "\n",
    "\n",
    "args = [\"--config\", \"examples/search-agent/local_1.5b_example.yaml\"]\n",
    "config, _ = load_expr_config(args, AgentRLConfig)\n",
    "config: AgentRLConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "### 配置环境变量\n",
    "\n",
    "我们预先分配SGLang服务器和PyTorch分布式启动的IP地址和端口，并设置相应的环境变量。\n",
    "\n",
    "这些环境变量会在引擎初始化时被读取。\n",
    "\n",
    "***在非notebook环境下，这些环境变量会被launcher设置，用户无需自行设置。***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from areal.utils.network import find_free_ports\n",
    "\n",
    "SGLANG_PORT, MASTER_PORT = 11451, 14514\n",
    "\n",
    "SGLANG_HOST = \"127.0.0.1\"\n",
    "\n",
    "# ----------------------------------------------------------------------------\n",
    "# Environment variables used by inference/train engines\n",
    "import os\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "os.environ[\"AREAL_LLM_SERVER_ADDRS\"] = f\"{SGLANG_HOST}:{SGLANG_PORT}\"\n",
    "os.environ[\"MASTER_ADDR\"] = \"127.0.0.1\"\n",
    "os.environ[\"MASTER_PORT\"] = str(MASTER_PORT)\n",
    "os.environ[\"RANK\"] = str(0)\n",
    "os.environ[\"WORLD_SIZE\"] = str(1)\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n",
    "os.environ[\"LOCAL_RANK\"] = str(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "### 启动SGLang服务器\n",
    "\n",
    "AReaL默认采用训推分离式架构，推理和训练异步执行，能够打满GPU资源、快速完成端到端训练。\n",
    "\n",
    "在这个样例中，强化学习的算法编排（GRPO）运行在GPU 0上。\n",
    "\n",
    "GPU 1运行一个推理服务器，强化学习的算法编排可以向GPU 1上的推理服务发送生成请求。\n",
    "\n",
    "下面的代码块在GPU 1上启动对应的推理服务。\n",
    "\n",
    "本次教程中使用`Qwen/Qwen2.5-1.5B`作为例子。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 启动sglang server\n",
    "from areal.api.cli_args import SGLangConfig\n",
    "from areal.utils.network import find_free_ports\n",
    "\n",
    "config.sglang.log_level = \"info\"\n",
    "config.sglang.decode_log_interval = 10\n",
    "sglang_cmd = SGLangConfig.build_cmd(\n",
    "    config.sglang,\n",
    "    tp_size=1,\n",
    "    base_gpu_id=1,\n",
    "    host=SGLANG_HOST,\n",
    "    port=SGLANG_PORT,\n",
    ")\n",
    "sglang_process = subprocess.Popen(\n",
    "    sglang_cmd,\n",
    "    stdout=sys.stdout,\n",
    "    stderr=sys.stderr,\n",
    ")\n",
    "\n",
    "print(\"sglang process is launched\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### 加载训练数据集\n",
    "\n",
    "使用HuggingFace `datasets` 包加载训练数据集，并查看数据集格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load search dataset\n",
    "from datasets import load_dataset\n",
    "\n",
    "print(\"dataset is at {}\".format(config.train_dataset.path))\n",
    "dataset = load_dataset(\n",
    "    path=\"json\",\n",
    "    split=\"train\",\n",
    "    data_files=config.train_dataset.path,\n",
    ")\n",
    "print(f\">>> dataset column names: {dataset.column_names}\")\n",
    "print(f\">>> example data: {dataset[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "### 导入必要的包和模块\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import json\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import uuid\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.distributed as dist\n",
    "from datasets import load_dataset\n",
    "from datasets.distributed import split_dataset_by_node\n",
    "from transformers import AutoTokenizer, PreTrainedTokenizerFast\n",
    "\n",
    "from areal.api.cli_args import (\n",
    "    GenerationHyperparameters,\n",
    "    load_expr_config,\n",
    ")\n",
    "from areal.api.engine_api import InferenceEngine\n",
    "from areal.api.io_struct import (\n",
    "    AllocationMode,\n",
    "    FinetuneSpec,\n",
    "    ModelRequest,\n",
    "    WeightUpdateMeta,\n",
    ")\n",
    "from areal.engine.ppo.actor import FSDPPPOActor\n",
    "from areal.engine.sglang_remote import RemoteSGLangEngine\n",
    "from areal.utils.data import concat_padded_tensors, tensor_container_to\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path)\n",
    "if tokenizer.pad_token_id is None:\n",
    "    tokenizer.pad_token_id = tokenizer.eos_token_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "使用`torchdata.stateful_dataloader.StatefulDataLoader` 作为dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup dataloader\n",
    "\n",
    "from torchdata.stateful_dataloader import StatefulDataLoader\n",
    "\n",
    "dataloader = StatefulDataLoader(\n",
    "    dataset,\n",
    "    batch_size=config.train_dataset.batch_size,\n",
    "    shuffle=True,\n",
    "    collate_fn=lambda x: x,\n",
    "    drop_last=True,\n",
    ")\n",
    "\n",
    "from itertools import cycle\n",
    "\n",
    "data_generator = cycle(dataloader)\n",
    "\n",
    "ft_spec = FinetuneSpec(\n",
    "    total_train_epochs=config.total_train_epochs,\n",
    "    dataset_size=len(dataloader) * config.train_dataset.batch_size,\n",
    "    train_batch_size=config.train_dataset.batch_size,\n",
    ")\n",
    "\n",
    "batch = next(data_generator)\n",
    "print(f\">>> The type of a batch is: {type(batch)}\\n\")\n",
    "print(f\">>> Each piece of data has keys: {batch[0].keys()}\\n\")\n",
    "print(f\">>> Example input question: {batch[0]['question']}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### 配置搜索工具"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "本地RAG服务器部署方式请见：[ASearcher 仓库](https://github.com/inclusionAI/ASearcher/blob/main/docs/training.md#b-training-a-search-agent-with-local-knowledge-base) - Step 2.\n",
    "\n",
    "通过5001端口给本地RAG服务器发送查询，并接收结果。\n",
    "\n",
    "如下展示了一个搜索关键词\"China\"的例子。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup tool\n",
    "\n",
    "import asyncio\n",
    "import json\n",
    "\n",
    "import aiohttp\n",
    "\n",
    "TOOL_SERVER_ADDR = \"localhost:5001\"\n",
    "\n",
    "\n",
    "async def call_search_tool(**req_meta):\n",
    "    async with aiohttp.ClientSession() as session:\n",
    "        async with session.post(\n",
    "            f\"http://{TOOL_SERVER_ADDR}/retrieve\",\n",
    "            json=req_meta,\n",
    "            timeout=aiohttp.ClientTimeout(total=120, sock_connect=120),\n",
    "        ) as response:\n",
    "            response.raise_for_status()\n",
    "            res = await response.json()\n",
    "            return res[\"result\"]\n",
    "\n",
    "\n",
    "result = (await call_search_tool(queries=[\"China\"], topk=5, return_scores=False))[0]\n",
    "print(json.dumps(result, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## 定义简单的智能体工作流"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### 模型输出\n",
    "\n",
    "使用prompt控制模型输出，模型输出应遵循特定格式：\n",
    "- `<think></think>` 包含模型思考过程\n",
    "- `<search></search>` 包含给本地RAG服务器的查询\n",
    "- `<answer></answer>` 包含模型输出的答案\n",
    "\n",
    "此外，使用`<information></information>` 包含RAG服务器返回的查询内容。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROMPT_TEMPLATE = \"\"\"A conversation between User and Assistant. The user asks a question, and the Assistant answers it. The Assistant analyzes the given question and information in the mind, retains important relevant information, calls a search engine to find necessary information, accesses web pages with certain urls, and provides the user with the answer. The Assistant conducts search by <search> query </search> and the top search results will be returned between <information> and </information>. The reasoning processes are enclosed within <think> </think>. Finally, the Assistant provides answer inside <answer> and </answer>, i.e. <answer> answer here </answer>. If there are multiple queries, ensure all answers are enclosed within <answer> </answer>, seperated with comma.\n",
    "\n",
    "User:\n",
    "{question}\n",
    "\n",
    "Assistant:\n",
    "<think>\"\"\"\n",
    "\n",
    "batch = next(data_generator)\n",
    "prompt = PROMPT_TEMPLATE.format(question=batch[0][\"question\"])\n",
    "\n",
    "print(f\">>> PROMPT: {prompt}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "通过`RemoteSGlangEngine`向已经启动的SGLang服务器发送生成请求，测试上述prompt控制下的模型输出。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "asyncio.get_running_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize inference engine\n",
    "rollout_engine = RemoteSGLangEngine(config.rollout)\n",
    "rollout_engine.initialize()\n",
    "\n",
    "# generation config\n",
    "gconfig = GenerationHyperparameters(\n",
    "    max_new_tokens=512, stop=[\"</search>\", \"</answer>\", \"</access>\"]\n",
    ")\n",
    "\n",
    "# tokenize the prompt\n",
    "input_ids = tokenizer([prompt], add_special_tokens=False)[\"input_ids\"][0]\n",
    "req = ModelRequest(rid=uuid.uuid4().hex, input_ids=input_ids, gconfig=gconfig)\n",
    "\n",
    "# generate rollout with inference engine\n",
    "resp = await rollout_engine.agenerate(req)\n",
    "completion_str = tokenizer.decode(resp.output_tokens)\n",
    "\n",
    "# logging\n",
    "print(f\">>> prompt str: {tokenizer.decode(resp.input_tokens)}\")\n",
    "print(f\">>> generated: {tokenizer.decode(resp.output_tokens)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "#### 解析智能体工具调用\n",
    "\n",
    "定义`parse_search_query`函数从模型输出解析调用搜索工具的查询。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parse tool calling\n",
    "\n",
    "import re\n",
    "\n",
    "\n",
    "def parse_search_query(text):\n",
    "    pattern = r\"<search>(.*?)</search>\"\n",
    "    matches = re.findall(pattern, text, re.DOTALL)\n",
    "    if matches:\n",
    "        return matches[-1].strip()\n",
    "    return None\n",
    "\n",
    "\n",
    "test_tool_str = \"<think> I would like to search for AI.</think>\\n<search> Artificial Intelligence </search>\"\n",
    "print(\">>> input: \", test_tool_str)\n",
    "print(\">>> search query: \", parse_search_query(test_tool_str))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "在模型输出上测试工具调用解析函数`parse_search_query`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate rollout with inference engine\n",
    "resp = await rollout_engine.agenerate(req)\n",
    "completion_str = tokenizer.decode(resp.output_tokens)\n",
    "\n",
    "# logging\n",
    "print(f\">>> prompt str: {tokenizer.decode(resp.input_tokens)}\")\n",
    "print(f\">>> generated: {tokenizer.decode(resp.output_tokens)}\")\n",
    "print(f\">>> search query: {parse_search_query(completion_str)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "#### 解析智能体答案\n",
    "\n",
    "定义 `parse_answer` 函数从模型输出中解析模型答案。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parse answer\n",
    "\n",
    "\n",
    "def parse_answer(text):\n",
    "    pattern = r\"<answer>(.*?)</answer>\"\n",
    "    matches = re.findall(pattern, text, re.DOTALL)\n",
    "    if matches:\n",
    "        return matches[-1].strip()\n",
    "    return None\n",
    "\n",
    "\n",
    "test_answer_str = (\n",
    "    \"<think> I already found the answer! </think>\\n<answer> 1997 </answer>\"\n",
    ")\n",
    "print(\">>> input: \", test_answer_str)\n",
    "print(\">>> answer: \", parse_answer(test_answer_str))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {},
   "source": [
    "在模型输出上测试答案解析函数`parse_answer`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate rollout with inference engine\n",
    "resp = await rollout_engine.agenerate(req)\n",
    "completion_str = tokenizer.decode(resp.output_tokens)\n",
    "\n",
    "# logging\n",
    "print(f\">>> prompt str: {tokenizer.decode(resp.input_tokens)}\")\n",
    "print(f\">>> generated: {tokenizer.decode(resp.output_tokens)}\")\n",
    "print(f\">>> answer: {parse_answer(completion_str)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {},
   "source": [
    "### 奖励函数\n",
    "\n",
    "我们默认使用F1 score作为奖励函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# F1 reward\n",
    "\n",
    "\n",
    "def f1_score(pred_ans, gt):\n",
    "    # 预处理文本（此处为简化版本）\n",
    "    pred_ans = pred_ans.strip().lower()\n",
    "    gt = gt.strip().lower()\n",
    "\n",
    "    pred_tokens = set(pred_ans.split())\n",
    "    gt_tokens = set(gt.split())\n",
    "\n",
    "    if not gt_tokens or not pred_tokens:\n",
    "        return 0\n",
    "\n",
    "    # 计算共同的词数\n",
    "    common_tokens = pred_tokens & gt_tokens\n",
    "\n",
    "    # 计算精确率和召回率\n",
    "    precision = len(common_tokens) / len(pred_tokens) if pred_tokens else 0\n",
    "    recall = len(common_tokens) / len(gt_tokens) if gt_tokens else 0\n",
    "\n",
    "    # 计算F1分数\n",
    "    f1 = 0\n",
    "    if precision + recall > 0:\n",
    "        f1 = 2 * (precision * recall) / (precision + recall)\n",
    "\n",
    "    return f1\n",
    "\n",
    "\n",
    "print(\n",
    "    \"f1_score('James Bond', 'James Bond'): {:.2f}\".format(\n",
    "        f1_score(\"James Bond\", \"James Bond\")\n",
    "    )\n",
    ")\n",
    "print(\n",
    "    \"f1_score('James Smith', 'James Bond'): {:.2f}\".format(\n",
    "        f1_score(\"James Smith\", \"James Bond\")\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32",
   "metadata": {},
   "source": [
    "### 定义搜索智能体工作流\n",
    "\n",
    "实现搜索智能体 (Search Agent) 的工作流非常简单，从一个初始问题出发，在每一轮：\n",
    "1. 调用推理引擎进行生成，当生成到EOS、`</search>`、`</answer>`之一时停止生成\n",
    "2. 如果检测到搜索查询，调用搜索工具，并将搜索结果加入到历史中\n",
    "3. 如果检测到答案，计算奖励并退出循环\n",
    "\n",
    "最后将数据组合成训练需要的形式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Implement search agent workflow\n",
    "\n",
    "\n",
    "class SearchAgentWorkflow:\n",
    "    def __init__(self, gconfig, tokenizer, max_tokens, max_turns, verbose):\n",
    "        self.gconfig = gconfig\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_tokens = max_tokens\n",
    "        self.max_turns = max_turns\n",
    "        self.verbose = verbose\n",
    "\n",
    "    async def arun_episode(self, engine: InferenceEngine, data):\n",
    "        prompt = PROMPT_TEMPLATE.format(question=data[\"question\"])\n",
    "\n",
    "        # an unique trajectory rid to ensure all requests goes to the same sglang server\n",
    "        rid = uuid.uuid4().hex\n",
    "\n",
    "        # trajectory (input ids/logprobs/loss mask)\n",
    "        input_ids = self.tokenizer.encode(prompt, add_special_tokens=False)\n",
    "        logprobs = [0.0] * len(input_ids)\n",
    "        loss_mask = [0] * len(input_ids)\n",
    "\n",
    "        answer, reward = None, 0\n",
    "\n",
    "        num_turns = 0\n",
    "        while num_turns < self.max_turns and len(input_ids) < self.max_tokens:\n",
    "            num_turns += 1\n",
    "\n",
    "            # LLM Request\n",
    "            req = ModelRequest(\n",
    "                rid=rid,\n",
    "                input_ids=input_ids,\n",
    "                gconfig=self.gconfig.new(n_samples=1),\n",
    "            )\n",
    "            resp = await engine.agenerate(req)\n",
    "            completion_str = self.tokenizer.decode(resp.output_tokens)\n",
    "\n",
    "            input_ids += resp.output_tokens\n",
    "            input_ids += resp.output_tokens\n",
    "            logprobs += resp.output_logprobs\n",
    "            loss_mask += [1] * resp.output_len\n",
    "\n",
    "            # parse search query & trigger tool call\n",
    "            search_query = parse_search_query(completion_str)\n",
    "            if search_query:\n",
    "                search_results = (\n",
    "                    await call_search_tool(\n",
    "                        queries=[search_query], topk=3, return_scores=False\n",
    "                    )\n",
    "                )[0]\n",
    "                search_results_str = (\n",
    "                    \"\\n\\n<information>\\n\"\n",
    "                    + \"\\n\\n\".join(\n",
    "                        [\n",
    "                            '<p title=\"{}\">\\n{}\\n</p>'.format(\n",
    "                                r[\"wikipedia_title\"], r[\"contents\"]\n",
    "                            )\n",
    "                            for r in search_results\n",
    "                        ]\n",
    "                    )\n",
    "                    + \"\\n</information>\"\n",
    "                )\n",
    "\n",
    "                search_token_ids = self.tokenizer.encode(\n",
    "                    search_results_str, add_special_tokens=False\n",
    "                )\n",
    "                input_ids += search_token_ids\n",
    "                logprobs += [0.0] * len(search_token_ids)\n",
    "                loss_mask += [0] * len(search_token_ids)\n",
    "\n",
    "            # parse answer\n",
    "            answer = parse_answer(completion_str)\n",
    "            if answer:\n",
    "                reward = max([f1_score(answer, gt) for gt in data[\"answer\"]])\n",
    "                break\n",
    "\n",
    "            if input_ids[-1] in [\n",
    "                self.tokenizer.pad_token_id,\n",
    "                self.tokenizer.eos_token_id,\n",
    "            ]:\n",
    "                break\n",
    "\n",
    "        if self.verbose:\n",
    "            print(f\"[LOGGING] turns={num_turns} length={len(input_ids)}\")\n",
    "\n",
    "        res = dict(\n",
    "            input_ids=torch.tensor(input_ids),\n",
    "            logprobs=torch.tensor(logprobs),\n",
    "            loss_mask=torch.tensor(loss_mask),\n",
    "            rewards=torch.tensor(float(reward)),\n",
    "            attention_mask=torch.ones(len(input_ids), dtype=torch.bool),\n",
    "        )\n",
    "        res = {k: v.unsqueeze(0) for k, v in res.items()}\n",
    "        return res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34",
   "metadata": {},
   "source": [
    "#### 测试搜索智能体工作流\n",
    "\n",
    "1. 创建推理引擎；\n",
    "2. 创建工作流，设定`max_new_tokens`, `max_turns` 和 `max_tokens`；\n",
    "3. 将工作流传入推理引擎进行批量生成。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize inference engine\n",
    "rollout = RemoteSGLangEngine(config.rollout)\n",
    "rollout.initialize()\n",
    "\n",
    "# TODO: create workflow\n",
    "workflow = SearchAgentWorkflow(\n",
    "    gconfig=GenerationHyperparameters(\n",
    "        max_new_tokens=512, stop=[\"</answer>\", \"</search>\"]\n",
    "    ),\n",
    "    tokenizer=tokenizer,\n",
    "    max_tokens=4096,\n",
    "    max_turns=32,\n",
    "    verbose=True,\n",
    ")\n",
    "sample_data = next(data_generator)[:4]\n",
    "res = await asyncio.gather(\n",
    "    *[workflow.arun_episode(rollout, sample_data[i]) for i in range(4)]\n",
    ")\n",
    "res = concat_padded_tensors(res)\n",
    "print(res)\n",
    "\n",
    "rollout.destroy()\n",
    "\n",
    "# log the trajectories\n",
    "traj_lens = res[\"attention_mask\"].sum(dim=1).numpy().tolist()\n",
    "for i in range(4):\n",
    "    token_ids = res[\"input_ids\"][i, : traj_lens[i]]\n",
    "    print(f\">>> Trajectory {i} >>>\\n{tokenizer.decode(token_ids)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36",
   "metadata": {},
   "source": [
    "### 让智能体工作流对每个问题生成多条轨迹\n",
    "\n",
    "\n",
    "类似GRPO的算法需要针对每个问题生成一组多条轨迹。\n",
    "\n",
    "我们可以通过一个asyncio的并行技巧同时高效地生成多个轨迹。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Group generation for GRPO\n",
    "\n",
    "\n",
    "class GroupedSearchAgentWorkflow:\n",
    "    def __init__(self, gconfig, tokenizer, max_tokens, max_turns, group_size, verbose):\n",
    "        self.gconfig = gconfig\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_tokens = max_tokens\n",
    "        self.max_turns = max_turns\n",
    "        self.group_size = group_size\n",
    "        self.verbose = verbose\n",
    "\n",
    "    async def arun_episode(self, engine, data):\n",
    "        workflows = [\n",
    "            SearchAgentWorkflow(\n",
    "                self.gconfig.new(n_samples=1),\n",
    "                self.tokenizer,\n",
    "                self.max_tokens,\n",
    "                self.max_turns,\n",
    "                self.verbose,\n",
    "            )\n",
    "            for _ in range(self.group_size)\n",
    "        ]\n",
    "        tasks = [workflow.arun_episode(engine, data) for workflow in workflows]\n",
    "        results = await asyncio.gather(*tasks)\n",
    "        return concat_padded_tensors(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize inference engine\n",
    "rollout = RemoteSGLangEngine(config.rollout)\n",
    "rollout.initialize()\n",
    "try:\n",
    "    # TODO: create workflow\n",
    "    workflow = GroupedSearchAgentWorkflow(\n",
    "        gconfig=GenerationHyperparameters(\n",
    "            max_new_tokens=512, stop=[\"</answer>\", \"</search>\"]\n",
    "        ),\n",
    "        tokenizer=tokenizer,\n",
    "        max_tokens=4096,\n",
    "        max_turns=32,\n",
    "        group_size=4,\n",
    "        verbose=True,\n",
    "    )\n",
    "    sample_data = next(data_generator)[:2]\n",
    "    res = rollout.rollout_batch(sample_data, workflow=workflow)\n",
    "    print(res)\n",
    "finally:\n",
    "    rollout.destroy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39",
   "metadata": {},
   "source": [
    "## 将智能体工作流接入强化学习训练流程\n",
    "\n",
    "上面我们已经测试好了负责的推理工作流，接下来我们需要将这个工作流接入到训练过程中。\n",
    "\n",
    "这需要我们额外创建一个专门针对PPO的训练引擎，并在training loop中循环调用推理和训练。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training for 5 steps\n",
    "\n",
    "workflow = GroupedSearchAgentWorkflow(\n",
    "    gconfig=GenerationHyperparameters(\n",
    "        max_new_tokens=512, stop=[\"</answer>\", \"</search>\"]\n",
    "    ),\n",
    "    tokenizer=tokenizer,\n",
    "    max_tokens=4096,\n",
    "    max_turns=32,\n",
    "    group_size=4,\n",
    "    verbose=True,\n",
    ")\n",
    "actor = FSDPPPOActor(config=config.actor)\n",
    "actor.initialize(None, ft_spec)\n",
    "\n",
    "rollout = RemoteSGLangEngine(config.rollout)\n",
    "rollout.initialize()\n",
    "\n",
    "weight_update_meta = WeightUpdateMeta.from_fsdp_xccl(\n",
    "    AllocationMode.from_str(\"sglang:d1p1t1+d1p1t1\")\n",
    ")\n",
    "\n",
    "warmup_steps = 1\n",
    "times = []\n",
    "for global_step in range(5):\n",
    "    if global_step >= warmup_steps:\n",
    "        tik = time.perf_counter()\n",
    "    batch = rollout.rollout_batch(next(data_generator), workflow=workflow)\n",
    "    print(batch)\n",
    "    batch = tensor_container_to(batch, actor.device)\n",
    "\n",
    "    logp = actor.compute_logp(batch)\n",
    "    batch[\"prox_logp\"] = logp\n",
    "\n",
    "    actor.compute_advantages(batch)\n",
    "\n",
    "    actor.ppo_update(batch)\n",
    "    actor.step_lr_scheduler()\n",
    "\n",
    "    rollout.pause()\n",
    "    actor.update_weights(weight_update_meta)\n",
    "    rollout.resume()\n",
    "\n",
    "    actor.set_version(global_step + 1)\n",
    "    rollout.set_version(global_step + 1)\n",
    "    if global_step >= warmup_steps:\n",
    "        times.append(time.perf_counter() - tik)\n",
    "print(times)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AReaL",
   "language": "python",
   "name": "areal"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
